#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: 邵奈一
@Email: shaonaiyi@163.com
@Date: 2024/11/12
@微信：shaonaiyi888
@微信公众号: 邵奈一 
"""
# 代码 3-28
# 加载数据并预处理
import numpy as np
import pandas as pd
import torch
from torch.autograd import Variable
from torch import nn

# 读取CSV文件，只选择第一列数据
data_csv = pd.read_csv('data3/data.csv', usecols=[1])
# 删除数据中的空值
data_csv = data_csv.dropna()
# 获取数据集的值
dataset = data_csv.values
# 将数据集的值转换为浮点型
dataset = dataset.astype('float32')
# 获取数据集中的最大值
max_value = np.max(dataset)
# 获取数据集中的最小值
min_value = np.min(dataset)
# 计算最大值和最小值的差值
scalar = max_value - min_value
# 将数据集中的每个值除以最大值和最小值的差值，然后将结果转换为列表
dataset = list(map(lambda x: x / scalar, dataset))

# 代码 3-29
# 创建数据集
def create_dataset(dataset, look_back=2):
    """
    创建数据集

    参数:
    dataset: 原始数据集
    look_back: 用于预测的回溯时间步数，默认为2

    返回值:
    dataX: 用于训练的输入数据
    dataY: 用于预测的目标数据
    """
    # 初始化用于训练的输入数据和目标数据
    dataX, dataY = [], []
    # 遍历数据集，每次取look_back长度的数据作为输入，look_back长度的数据的下一个值作为目标
    for i in range(len(dataset) - look_back):
        # 从数据集中取出look_back长度的数据
        a = dataset[i: (i + look_back)]
        # 将取出的数据添加到输入数据中
        dataX.append(a)
        # 将目标数据添加到目标数据中
        dataY.append(dataset[i + look_back])
    # 将输入数据和目标数据转换为numpy数组并返回
    return np.array(dataX), np.array(dataY)

# 创建好输入输出
data_X, data_Y = create_dataset(dataset)
# 划分训练集和测试集，70% 作为训练集
train_size = int(len(data_X) * 0.7)
test_size = len(data_X) - train_size
train_X = data_X[: train_size]
train_Y = data_Y[: train_size]
test_X = data_X[train_size:]
test_Y = data_Y[train_size:]

# 代码3-30
# 改变数据形状成(seq, batch, feature)
train_X = train_X.reshape(-1, 1, 2) # -1会根据其他维度的大小自动推断出合适的值
train_Y = train_Y.reshape(-1, 1, 1)
test_X = test_X.reshape(-1, 1, 2)

train_x = torch.from_numpy(train_X)
train_y = torch.from_numpy(train_Y)
test_x = torch.from_numpy(test_X)

# 代码3-31
# 构建LSTM网络
class lstm_reg(nn.Module):

    def __init__(self, input_size, hidden_size, output_size=1, num_layers=2):
        super(lstm_reg, self).__init__()
        self.rnn = nn.LSTM(input_size, hidden_size, num_layers)  # rnn
        self.reg = nn.Linear(hidden_size, output_size)  # 回归

    def forward(self, x):
        x, _ = self.rnn(x)  # (seq, batch, hidden)
        s, b, h = x.shape
        x = x.view(s * b, h)  # 转换成线性层的输入格式
        x = self.reg(x) # 全连接层，用于将LSTM的输出转换为最终的回归输出
        x = x.view(s, b, -1)
        return x

# 代码3-32
# 构建RNN网络
net = lstm_reg(2, 4)
criterion = nn.MSELoss()  # 交叉熵损失函数
optimizer = torch.optim.Adam(net.parameters(), lr=1e-2)  # Adam优化算法

# 代码3-33
# 训练网络
for e in range(1000):
    var_x = Variable(train_x)
    var_y = Variable(train_y)
    # 前向传播
    out = net(var_x)
    loss = criterion(out, var_y)
    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (e) % 100 == 0:  # 每100次输出结果
        print('Epoch: {}, Loss: {:.5f}'.format(e, loss))

# 代码3-34
# 性能评估
# 用训练好的模型去预测结果
import matplotlib.pyplot as plt
# 设置字体
plt.rcParams['font.family'] = 'Heiti TC'
net = net.eval()
data_X = data_X.reshape(-1, 1, 2)
data_X = torch.from_numpy(data_X)
var_data = Variable(data_X)
pred_test = net(var_data)  # 测试集的预测结果
# 改变输出的格式
pred_test = pred_test.view(-1).data.numpy()
# 画出实际结果和预测的结果
plt.plot(pred_test, 'r', label='预测值')
plt.plot(dataset, 'b', label='真实值')
plt.legend(['预测值', '真实值'], loc='upper left')
# plt.legend(loc='best') # 不好使
plt.savefig('tmp/lstm-time.png', dpi=1080)
